#include "task.h"
#include "klib.h"
#include "os_cfg.h"
#include "cpu.h"
#include "log.h"
#include "cpu_instr.h"
#include "irq.h"
#include "mem.h"
#include "mmu.h"
#include "syscall.h"
#include "fs.h"
#include "elf.h"
#include "file.h"

static task_manager_t task_manager;
static uint32_t       idle_task_stack[IDLE_TASK_STACK_SIZE];
static task_t         task_table[TASK_NR];
static mutex_t        task_table_mutex;

static int tss_init(task_t *task, int flag, uint32_t entry, uint32_t esp)
{
    // allocate tss selector and save it in gdt_table
    int tss_selector = gdt_alloc_desc();
    if (tss_selector < 0)
    {
        log_printf("alloc tss failed!");
        return -1;
    }
    segment_desc_set(tss_selector, (uint32_t)&task->tss, sizeof(tss_t), SEG_P | SEG_DPL0 | SEG_TYPE_TSS);
    // tss init in task
    kernel_memset(&task->tss, 0, sizeof(tss_t));

    uint32_t kernel_stack = mem_alloc_page();    // for privilege 0
    if (kernel_stack == 0)
    {
        goto tss_init_failed;
    }
    int code_sel, data_sel;
    if (flag & TASK_FLAGS_SYSTEM)
    {
        code_sel = KERNEL_SELECTOR_CS;
        data_sel = KERNEL_SELECTOR_DS;
    }
    else
    {
        code_sel = task_manager.app_code_sel | SEG_CPL3;
        data_sel = task_manager.app_data_sel | SEG_CPL3;
    }

    task->tss.eip  = entry;
    task->tss.esp  = esp;
    task->tss.esp0 = kernel_stack + MEM_PAGE_SIZE;
    task->tss.ss   = data_sel;
    task->tss.ss0  = KERNEL_SELECTOR_DS;
    task->tss.es = task->tss.ds = task->tss.fs = task->tss.gs = data_sel;
    task->tss.cs                                              = code_sel;
    task->tss.eflags                                          = EFLAGS_DEFAULT | EFLAGS_IF;
    task->tss_selector                                        = tss_selector;
    uint32_t page_dir                                         = mem_create_uvm();
    if (page_dir == 0)
    {
        goto tss_init_failed;
    }
    task->tss.cr3 = page_dir;
    return 0;
tss_init_failed:
    gdt_free_selector(tss_selector);
    if (kernel_stack)
    {
        mem_free_page(kernel_stack);
    }
    return -1;
}
int task_init(task_t *task, const char *name, int flag, uint32_t entry, uint32_t esp)
{
    ASSERT(task != (task_t *)0);
    tss_init(task, flag, entry, esp);

    kernel_strncpy(task->name, name, TASK_NAME_SIZE);
    task->state       = TASK_CREATED;
    task->sleep_ticks = 0;
    task->time_ticks  = TASK_TIME_SLICE_DEFAULT;
    task->slice_ticks = task->time_ticks;

    list_node_init(&task->run_node);
    list_node_init(&task->all_node);
    list_node_init(&task->wait_node);

    kernel_memset(task->file_table, 0, sizeof(task->file_table));
    // kernel_memset(task->file_table,0,sizeof(filie_t*)*TASK_OFILE_NR);

    irq_state_t irq_state = irq_enter_protect();
    task->pid             = (uint32_t)&task;
    task->parent          = (task_t *)0;
    task->heap_start      = 0;
    task->heap_end        = 0;
    task->status          = 0;
    list_insert_last(&task_manager.task_list, &task->all_node);
    log_printf("task init:task_name:%s, task_id: %x, task->tss.cr3:%x", task->name, task->pid, task->tss.cr3);
    irq_exit_protect(irq_state);
    // uint32_t *pesp=(uint32_t*)esp;
    //    if(pesp){
    //          *(--pesp) = entry;
    //          *(--pesp)=0;
    //          *(--pesp)=0;
    //          *(--pesp)=0;
    //          *(--pesp)=0;
    //          task->stack=pesp;
    //    }
    return 0;
}
void task_start(task_t *task)
{
    irq_state_t irq_state = irq_enter_protect();
    task_set_ready(task);
    irq_exit_protect(irq_state);
}
// implement on start.S
void simple_switch(uint32_t **from, uint32_t *to);
void task_switch_from_to(task_t *from, task_t *to)
{
    switch_to_tss(to->tss_selector);
    // simple_switch(&from->stack,to->stack);
}

void task_first_init(void)
{
    void first_task_entry(void);
    // copy first task
    extern uint8_t s_first_task[], e_first_task[];    // from kernel.ld
    uint32_t       first_start = (uint32_t)first_task_entry;
    uint32_t       copy_size   = (uint32_t)(e_first_task - s_first_task);
    uint32_t       alloc_size =
        10 * MEM_PAGE_SIZE;    // it is enough for first task, because only init word needs to do ,no app
    ASSERT(copy_size < alloc_size);

    task_t *first_task = task_first_task();
    task_init(first_task, "first task", 0, first_start, first_start + alloc_size);
    ltr(first_task->tss_selector);

    task_manager.current_task          = first_task;
    task_manager.first_task.heap_start = (uint32_t)e_first_task;
    task_manager.first_task.heap_end   = (uint32_t)e_first_task;
    mmu_set_page_dir(task_manager.first_task.tss.cr3);

    mem_alloc_page_for(first_start, alloc_size, PTE_P | PTE_W | PTE_U);
    kernel_memcpy((void *)first_start, s_first_task, copy_size);
    task_start(&task_manager.first_task);
}
task_t *task_first_task(void)
{
    return &task_manager.first_task;
}
static void idel_task_entry(void)
{
    for (;;)
    {
        hlt();
    }
}
task_t *task_current_task(void)
{
    return task_manager.current_task;
}
void task_manager_init(void)
{
    kernel_memset((void *)task_table, 0, sizeof(task_table));
    mutex_init(&task_table_mutex);

    int sel = gdt_alloc_desc();
    segment_desc_set(sel,
                     0x00000000,
                     0xffffffff,
                     SEG_P | SEG_DPL3 | SEG_S_CODE_DATA | SEG_TYPE_DATA | SEG_TYPE_DATA_WRITE | SEG_D);
    task_manager.app_data_sel = sel;
    sel                       = gdt_alloc_desc();
    segment_desc_set(sel,
                     0x00000000,
                     0xffffffff,
                     SEG_P | SEG_DPL3 | SEG_S_CODE_DATA | SEG_TYPE_CODE | SEG_TYPE_DATA_WRITE | SEG_D);
    task_manager.app_code_sel = sel;

    list_init(&task_manager.task_list);
    list_init(&task_manager.ready_list);
    list_init(&task_manager.sleep_list);
    task_manager.current_task = (task_t *)0;
    task_init(&task_manager.idle_task,
              "idle task",
              TASK_FLAGS_SYSTEM,
              (uint32_t)idel_task_entry,
              (uint32_t)(idle_task_stack + IDLE_TASK_STACK_SIZE));
    task_start(&task_manager.idle_task);
}
void task_set_ready(task_t *task)
{
    ASSERT(task != (task_t *)0);
    if (task == &task_manager.idle_task)
    {
        return;
    }
    list_insert_last(&task_manager.ready_list, &task->run_node);
    task->state = TASK_READY;
}
void task_set_block(task_t *task)
{
    ASSERT(task != (task_t *)0);
    if (task == &task_manager.idle_task)
    {
        return;
    }
    list_remove(&task_manager.ready_list, &task->run_node);
}
int sys_yield(void)
{
    if (list_count(&task_manager.ready_list) <= 1)
    {
        return -1;
    }
    irq_state_t irq_state = irq_enter_protect();
    if (list_count(&task_manager.ready_list) > 1)
    {
        task_t *curr_task = task_current_task();
        task_set_block(curr_task);
        task_set_ready(curr_task);
        task_dispatch();
    }
    irq_exit_protect(irq_state);
    return 0;
}
void sys_exit(int status)
{
    task_t *curr_task = task_current_task();
    for (int fd = 0; fd < TASK_OFILE_NR; fd++)
    {
        file_t *file = curr_task->file_table[fd];
        if (file)
        {
            sys_close(fd);
            curr_task->file_table[fd] = (file_t *)0;
        }
    }
    // set the child parent to first_task
    int zombi_child_task_exist_flag = 0;
    mutex_lock(&task_table_mutex);
    for (int i = 0; i < TASK_NR; i++)
    {
        task_t *task = task_table + i;
        if (task->parent == curr_task)
        {
            task->parent = &task_manager.first_task;
            if (task->state == TASK_ZOMBIE)
            {
                zombi_child_task_exist_flag = 1;
            }
        }
    }
    mutex_unlock(&task_table_mutex);
    irq_state_t irq_state = irq_enter_protect();
    task_t     *parent    = curr_task->parent;
    if (zombi_child_task_exist_flag && (parent != &task_manager.first_task))
    {
        if (task_manager.first_task.state == TASK_WAITING)
        {
            task_set_ready(&task_manager.first_task);
        }
    }
    // if the parent was waiting the child exit, parent should wait
    if (curr_task->parent->state == TASK_WAITING)
    {
        task_set_ready(curr_task->parent);
    }
    curr_task->status = status;
    curr_task->state  = TASK_ZOMBIE;
    task_set_block(curr_task);
    task_dispatch();
    irq_exit_protect(irq_state);
}
int sys_wait(int *status)
{
    task_t *curr_task = task_current_task();
    for (;;)
    {
        mutex_lock(&task_table_mutex);
        for (int i = 0; i < TASK_NR; i++)
        {
            task_t *task = task_table + i;
            if (task->parent != curr_task)
            {
                continue;
            }
            if (task->state == TASK_ZOMBIE)
            {
                int pid = task->pid;
                *status = task->status;
                mem_destroy_uvm(task->tss.cr3);
                mem_free_page(task->tss.esp0 - MEM_PAGE_SIZE);
                kernel_memset(task, 0, sizeof(task_t));
                mutex_unlock(&task_table_mutex);
                return pid;
            }
        }
        mutex_unlock(&task_table_mutex);
        irq_state_t irq_state = irq_enter_protect();
        task_set_block(curr_task);
        curr_task->state = TASK_WAITING;
        task_dispatch();
        irq_exit_protect(irq_state);
    }
    return 0;
}
task_t *task_next_task(void)
{
    if (list_empty(&task_manager.ready_list))
    {
        return &task_manager.idle_task;
    }
    list_node_t *first_node = list_first(&task_manager.ready_list);
    return container_of(first_node, task_t, run_node);
}
void task_dispatch(void)
{
    // ASSERT(!list_empty(&task_manager.ready_list));
    irq_state_t irq_state = irq_enter_protect();
    task_t     *to        = task_next_task();
    if (to != task_current_task())
    {
        task_t *from              = task_current_task();
        task_manager.current_task = to;
        to->state                 = TASK_RUNNING;
        task_switch_from_to(from, to);
    }
    irq_exit_protect(irq_state);
}
void task_time_tick(void)
{
    task_t *current_task = task_current_task();
    if (--current_task->slice_ticks == 0)
    {
        current_task->slice_ticks = current_task->time_ticks;
        task_set_block(current_task);
        task_set_ready(current_task);
        task_dispatch();
    }
    list_node_t *curr = list_first(&task_manager.sleep_list);
    while (curr)
    {
        list_node_t *next = curr->next;
        task_t      *task = container_of(curr, task_t, run_node);
        if (--task->sleep_ticks == 0)
        {
            task_set_wakeup(task);
            task_set_ready(task);
        }
        curr = next;
    }
    task_dispatch();
}

void task_set_sleep(task_t *task, uint32_t ticks)
{
    ASSERT(task != (task_t *)0);
    if (ticks == 0)
    {
        return;
    }
    task->sleep_ticks = ticks;
    task->state       = TASK_SLEEP;
    list_insert_last(&task_manager.sleep_list, &task->run_node);
}
void task_set_wakeup(task_t *task)
{
    list_remove(&task_manager.sleep_list, &task->run_node);
}
void sys_sleep(uint32_t ms)
{
    irq_state_t irq_state    = irq_enter_protect();
    task_t     *current_task = task_current_task();
    task_set_block(current_task);
    task_set_sleep(current_task, (ms + (OS_TICKS_MS - 1)) / OS_TICKS_MS);
    task_dispatch();
    irq_exit_protect(irq_state);
}
int sys_getpid()
{
    return task_current_task()->pid;
}
void task_uninit(task_t *task)
{
    if (task->tss_selector)
    {
        gdt_free_selector(task->tss_selector);
    }
    if (task->tss.esp0)
    {
        mem_free_page(task->tss.esp0 - MEM_PAGE_SIZE);
    }
    if (task->tss.cr3)
    {
        mem_destroy_uvm(task->tss.cr3);
    }
    kernel_memset(task, 0, sizeof(task_t));
}
static task_t *alloc_task()
{
    task_t *task = (task_t *)0;
    mutex_lock(&task_table_mutex);
    for (int i = 0; i < TASK_NR; i++)
    {
        task_t *curr = task_table + i;
        if (curr->name[0] == '\0')
        {
            task = curr;
            break;
        }
    }
    mutex_unlock(&task_table_mutex);
    return task;
}
static void free_task(task_t *task)
{
    mutex_lock(&task_table_mutex);
    task->name[0] = '\0';
    mutex_unlock(&task_table_mutex);
}
static void copy_opend_files(task_t *child_task)
{
    task_t *parent = task_current_task();
    for (int i = 0; i < TASK_OFILE_NR; i++)
    {
        file_t *file = parent->file_table[i];
        if (file)
        {
            file_inc_ref(file);
            child_task->file_table[i] = file;
        }
    }
}
int sys_fork()
{
    int     ret         = -1;
    task_t *parent_task = task_current_task();
    task_t *child_task  = alloc_task();
    if (child_task == (task_t *)0)
    {
        goto fork_failed;
    }
    syscall_frame_t *frame = (syscall_frame_t *)(parent_task->tss.esp0 - sizeof(syscall_frame_t));
    int              err   = task_init(child_task,
                        parent_task->name,
                        0,
                        frame->eip,
                        frame->esp + sizeof(uint32_t) * SYSCALL_GATE_PARAM_COUNT);
    if (err < 0)
    {
        goto fork_failed;
    }
    copy_opend_files(child_task);
    tss_t *tss = &child_task->tss;
    tss->eax   = 0;
    tss->edx   = frame->edx;
    tss->ebx   = frame->ebx;
    tss->ecx   = frame->ecx;
    tss->esi   = frame->esi;
    tss->edi   = frame->edi;
    tss->ebp   = frame->ebp;

    tss->cs            = frame->cs;
    tss->ds            = frame->ds;
    tss->es            = frame->es;
    tss->fs            = frame->fs;
    tss->gs            = frame->gs;
    tss->eflags        = frame->eflags;
    child_task->parent = parent_task;
    if ((child_task->tss.cr3 = mem_copy_uvm(parent_task->tss.cr3)) < 0)
    {
        goto fork_failed;
    }
    task_start(child_task);

    // tss->cr3=parent_task->tss.cr3;
    return child_task->pid;
fork_failed:
    if (child_task)
    {
        task_uninit(child_task);
        free_task(child_task);
    }
    return ret;
}
static int load_phdr(int fd, Elf32_Phdr *phdr, uint32_t page_dir)
{
    int err = mem_alloc_for_page_dir(page_dir, phdr->p_vaddr, phdr->p_memsz, PTE_P | PTE_U | PTE_W);
    if (err < 0)
    {
        log_printf("load phdr failed no memory");
        return -1;
    }
    if (sys_lseek(fd, phdr->p_offset, 0) < 0)
    {
        log_printf("load_phdr, lseek failed");
        return -1;
    }
    uint32_t vaddr = phdr->p_vaddr;
    uint32_t size  = phdr->p_filesz;
    while (size > 0)
    {
        int      curr_size = size > MEM_PAGE_SIZE ? MEM_PAGE_SIZE : size;
        uint32_t paddr     = mem_get_paddr(page_dir, vaddr);
        if (sys_read(fd, (char *)paddr, curr_size) < curr_size)
        {
            log_printf("load_phdr, read failed");
            return -1;
        }
        size -= curr_size;
        vaddr += curr_size;
    }
}

static uint32_t load_elf_file(task_t *task, const char *name, uint32_t page_dir)
{
    Elf32_Ehdr elf_hdr;
    Elf32_Phdr elf_phdr;
    int        fd = sys_open(name, 0);
    if (fd < 0)
    {
        log_printf("open file failed:%s,error_no:%d", name, fd);
        goto load_failed;
    }
    int cnt = sys_read(fd, (char *)&elf_hdr, sizeof(elf_hdr));
    if (cnt < sizeof(Elf32_Ehdr))
    {
        log_printf("elf hdr too small.size=%d", cnt);
        goto load_failed;
    }
    if ((elf_hdr.e_ident[0] != 0x7F) || (elf_hdr.e_ident[1] != 'E') || (elf_hdr.e_ident[2] != 'L') ||
        (elf_hdr.e_ident[3] != 'F'))
    {
        log_printf("check elf ident failed");
        goto load_failed;
    }
    uint32_t e_phoff = elf_hdr.e_phoff;
    for (int i = 0; i < elf_hdr.e_phnum; i++, e_phoff += elf_hdr.e_phentsize)
    {
        if (sys_lseek(fd, e_phoff, 0) < 0)
        {
            log_printf("lseek elf program header failed");
            goto load_failed;
        }
        cnt = sys_read(fd, (char *)&elf_phdr, sizeof(elf_phdr));
        if (cnt < sizeof(elf_phdr))
        {
            log_printf("read program header failed");
            goto load_failed;
        }
        if (elf_phdr.p_type != PT_LOAD || (elf_phdr.p_vaddr < MEM_TASK_BASE))
        {
            continue;
        }
        int err = load_phdr(fd, &elf_phdr, page_dir);
        if (err < 0)
        {
            log_printf("load program header failed");
            goto load_failed;
        }
        task->heap_start = elf_phdr.p_vaddr + elf_phdr.p_memsz;
        task->heap_end   = task->heap_start;
    }
    sys_close(fd);
    return elf_hdr.e_entry;
load_failed:
    if (fd)
    {
        sys_close(fd);
    }
    return 0;
}
static int copy_args(char *to, uint32_t page_dir, char **argv, char **env)
{
    int argc = strings_count(argv);
    int envc = strings_count(env);
    // task_args
    task_args_t task_args;
    task_args.argc = argc;    // 1+(argv==0?0:1);//+(env==0?0:1);
    task_args.argv = (char **)(to + sizeof(task_args_t));
    task_args.env  = (char **)(to + sizeof(task_args_t) + sizeof(char *) * (argc + 1));
    // args arr
    char **dest_argv_tb = (char **)mem_get_paddr(page_dir, (uint32_t)(to + sizeof(task_args_t)));
    char **dest_env_tb =
        (char **)mem_get_paddr(page_dir, (uint32_t)(to + sizeof(task_args_t) + sizeof(char *) * (argc + 1)));
    // args
    char *dest_arg = to + sizeof(task_args_t) + sizeof(char *) * (argc + 1) + sizeof(char *) * (envc + 1);
    ASSERT(dest_argv_tb != 0);
    ASSERT(dest_env_tb != 0);
    for (int i = 0; i < argc; i++)
    {
        char *from = argv[i];
        int   len  = kernel_strlen(from) + 1;    // plus '\0'
        int   err  = mem_copy_uvm_data((uint32_t)dest_arg, page_dir, (uint32_t)from, len * sizeof(char));
        ASSERT(err >= 0);
        dest_argv_tb[i] = dest_arg;
        dest_arg += len * sizeof(char);
    }
    // env
    for (int i = 0; i < envc; i++)
    {
        char *from = env[i];
        int   len  = kernel_strlen(from) + 1;    // plus '\0'
        int   err  = mem_copy_uvm_data((uint32_t)dest_arg, page_dir, (uint32_t)from, len * sizeof(char));
        ASSERT(err >= 0);
        dest_env_tb[i] = dest_arg;
        dest_arg += len * sizeof(char);
    }
    if (argc)
    {
        dest_argv_tb[argc] = '\0';
    }
    if (envc)
    {
        dest_env_tb[envc] = '\0';
    }
    return mem_copy_uvm_data((uint32_t)to, page_dir, (uint32_t)&task_args, sizeof(task_args_t));
}
/*static int copy_args(char * to,uint32_t page_dir,char **argv,char **env){
    int argc=strings_count(argv);
    task_args_t task_args;
    task_args.argc=argc;
    task_args.argv=(char**)(to+sizeof(task_args_t));
    char * dest_arg=to+sizeof(task_args_t)+sizeof(char*)*(argc+1);
    char **dest_argv_tb=(char**)mem_get_paddr(page_dir,(uint32_t)(to+sizeof(task_args_t)));
    ASSERT(dest_argv_tb != 0);
    for(int i=0;i<argc;i++){
        char * from=argv[i];
        int len=kernel_strlen(from)+1;//plus '\0'
        int err=mem_copy_uvm_data((uint32_t)dest_arg,page_dir,(uint32_t)from,len*sizeof(char));
        ASSERT(err>=0);
        dest_argv_tb[i]=dest_arg;
        dest_arg+=len*sizeof(char);
    }
    if(argc){
        dest_argv_tb[argc]='\0';
    }
    return mem_copy_uvm_data((uint32_t)to,page_dir,(uint32_t)&task_args,sizeof(task_args_t));
}*/

int sys_execve(char *name, char **argv, char **env)
{
    task_t *task = task_current_task();

    kernel_strncpy(task->name, get_file_name(name), TASK_NAME_SIZE);
    // create new pdt
    uint32_t new_page_dir = mem_create_uvm();
    // for safe, if failed ,we can go back and recover
    uint32_t old_page_dir = task->tss.cr3;
    if (!new_page_dir)
    {
        goto exec_failed;
    }
    // load elf file
    uint32_t entry = load_elf_file(task, name, new_page_dir);
    if (entry == 0)
    {
        goto exec_failed;
    }
    // set stack, this remained space is for argc and argv need transfer to the shell
    uint32_t stack_top = MEM_TASK_STACK_TOP - MEM_TASK_ARG_SIZE;
    int      err       = mem_alloc_for_page_dir(new_page_dir,
                                     MEM_TASK_STACK_TOP - MEM_TASK_STACK_SIZE,
                                     MEM_TASK_STACK_SIZE,
                                     PTE_P | PTE_W | PTE_U);
    if (err < 0)
    {
        log_printf("memory alloc stack space failed");
        goto exec_failed;
    }
    err = copy_args((char *)stack_top, new_page_dir, argv, env);
    if (err < 0)
    {
        goto exec_failed;
    }
    // deal with stack problem and set a clear space for new proc
    syscall_frame_t *frame = (syscall_frame_t *)(task->tss.esp0 - sizeof(syscall_frame_t));
    frame->eip             = entry;
    frame->eax = frame->ebx = frame->ecx = frame->edx = 0;
    frame->edi = frame->esi = frame->ebp = 0;
    frame->eflags                        = EFLAGS_IF | EFLAGS_DEFAULT;
    frame->esp                           = stack_top - sizeof(uint32_t) * SYSCALL_GATE_PARAM_COUNT;

    task->tss.cr3 = new_page_dir;
    mmu_set_page_dir(new_page_dir);
    mem_destroy_uvm(old_page_dir);
    return 0;
exec_failed:
    if (new_page_dir)
    {
        task->tss.cr3 = old_page_dir;
        mmu_set_page_dir(old_page_dir);
        mem_destroy_uvm(new_page_dir);
    }
    return -1;
}

int task_alloc_fd(file_t *file)
{
    int     fd   = -1;
    task_t *task = task_current_task();
    for (int i = 0; i < TASK_OFILE_NR; i++)
    {
        if (task->file_table[i] == (file_t *)0)
        {
            task->file_table[i] = file;
            return i;
        }
    }
    return fd;
}

void task_release_fd(int fd)
{
    if ((fd >= 0) && (fd < TASK_OFILE_NR))
    {
        task_current_task()->file_table[fd] = (file_t *)0;
    }
}
file_t *task_file(int fd)
{
    if ((fd >= 0) && (fd < TASK_OFILE_NR))
    {
        file_t *file = task_current_task()->file_table[fd];
        return file;
    }
    return (file_t *)0;
}